package afaas.framework.util;

import java.io.File;
import java.io.FileInputStream;
import java.net.URL;
import java.net.URLDecoder;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarInputStream;

public class PackageScanUtil {

    public static final String FILE_PROTOCOL = "file";
    public static final String JAR_PROTOCOL = "jar";
    public static final String UTF8_CHARSET = "UTF-8";
    public static final String CLASS_SUFFIX = ".class";
    public static final String CLASS_PATH_STUB = "classes";
    public static final int OFFSET_CLASS_PATH_STUB = 8;
    public static final String JAR_START_STUB = "file:";
    public static final String JAR_END_STUB = "!";

    @SuppressWarnings("rawtypes")
    public static Set<Class> scan(String packageName) throws RuntimeException{
        return scan(packageName, PackageScanUtil.class.getClassLoader(), UTF8_CHARSET);
    }

    @SuppressWarnings("rawtypes")
    public static Set<Class> scan(String packageName, ClassLoader classLoader, String charset) throws RuntimeException {

        CheckUtil.checkNotNull(packageName);

        String packageDirName = packageName.replace('.', '/');
        ClassLoader currentClassLoader = classLoader != null ? classLoader : Thread.currentThread().getContextClassLoader();

        Enumeration<URL> resources;
        try {
            resources = currentClassLoader.getResources(packageDirName);
            if (!resources.hasMoreElements()) {
                return new HashSet<>();
            }

            URL url = resources.nextElement();
            if(FILE_PROTOCOL.equals(url.getProtocol())){
                return scanDirectory(url, charset, classLoader);
            }else if(JAR_PROTOCOL.equals(url.getProtocol())){
                return scanJar(url, charset, classLoader);
            }else{
                throw new RuntimeException("Illegal file, url is " + url.toString());
            }

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @SuppressWarnings("rawtypes")
    private static Set<Class> scanDirectory(URL url , String charset, ClassLoader classloader) throws Exception{
        Set<Class> sets = new HashSet<>();

        File file = new File(URLDecoder.decode(url.getFile(), charset));
        if(!file.isDirectory()) {
            throw new RuntimeException("Illegal package, package is " + file.getAbsolutePath());
        }

        File[] files = file.listFiles();
        if(files == null){
            return new HashSet<>();
        }
        for(File fileE : files) {

            String absolutePath = fileE.getAbsolutePath();
            if(!absolutePath.endsWith(CLASS_SUFFIX)) {
                continue;
            }

            String classPath = absolutePath.substring(
                    absolutePath.lastIndexOf(CLASS_PATH_STUB) + OFFSET_CLASS_PATH_STUB,
                    absolutePath.indexOf(CLASS_SUFFIX)
            );

            String className = classPath.replaceAll("/", ".");
            Class clazz = Class.forName(className, Boolean.TRUE, classloader);
            sets.add(clazz);
        }

        return sets;

    }

    @SuppressWarnings("rawtypes")
    private static Set<Class> scanJar(URL url, String charset, ClassLoader classloader) throws Exception{
        Set<Class> sets = new HashSet<>();

        String urlFile = URLDecoder.decode(url.getFile(), charset);
        int jarStart = urlFile.indexOf(JAR_START_STUB) + JAR_START_STUB.length();
        int jarEnd = urlFile.indexOf(JAR_END_STUB);
        int packageStart = jarEnd + 2;
        int packageEnd = urlFile.length();

        String jarFileName = urlFile.substring(jarStart, jarEnd);
        String packageName = urlFile.substring(packageStart, packageEnd);

        JarInputStream jarInputStream = new JarInputStream(new FileInputStream(jarFileName));

        JarEntry jarEntry = jarInputStream.getNextJarEntry();
        while (jarEntry != null) {
            String jarEntryName = jarEntry.getName();
            if (jarEntryName.startsWith(packageName) && jarEntryName.endsWith(CLASS_SUFFIX)) {
                String className = jarEntryName.replaceAll("/", ".");
                className = className.substring(0, className.indexOf(CLASS_SUFFIX));
                Class clazz = Class.forName(className, Boolean.TRUE, classloader);
                sets.add(clazz);
            }
            jarEntry = jarInputStream.getNextJarEntry();
        }

        return sets;
    }

    public static void main(String[] args) {
        System.out.println(PackageScanUtil.scan("afaas.framework.util"));
    }
}

